import tensorflow as tf
import numpy as np


def _int64_list_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def write_to_file(output, offset = 0):
    writer = tf.python_io.TFRecordWriter(output)
    for i in range(100):
        size = np.random.randint(low = 0, high = 10)
        rnd_idx = np.random.randint(low = 0, high = 100, size = (size,))
        # Example -> Features -> Feature
        features = tf.train.Features(feature = {
            'index': _int64_feature(offset + i),
            'ids':  _int64_list_feature(rnd_idx)})
        print("ids = %r" % rnd_idx)
        example = tf.train.Example(features = features)
        writer.write(example.SerializeToString())

write_to_file("random_ids1.tfrecords")
write_to_file("random_ids2.tfrecords", 100)
